import torch

class LineSearch:
    def __init__(self, max_iters=100, c1=0.1, c2=0.9):
        """
        初始化非精确线搜索类

        参数：
        - max_iters：最大迭代次数，默认为100
        - c1：Armijo 准则中的参数，默认为0.1
        - c2：Wolfe 准则中的参数，默认为0.9
        """
        self.max_iters = max_iters
        self.c1 = c1
        self.c2 = c2

    def armijo(self, f, grad, x, d, alpha0=1.0):
        """
        Armijo 准则

        参数：
        - f：目标函数
        - grad：目标函数的梯度
        - x：当前点
        - d：搜索方向
        - alpha0：初始步长，默认为1.0

        返回：
        - alpha：满足 Armijo 准则的步长
        """
        alpha = alpha0
        # print("alpha1",alpha)
        print("alpha1",alpha)
        for _ in range(self.max_iters):
            if f(x + alpha * d) <= f(x) + self.c1 * alpha * torch.dot(grad, d):
                print("alpha2",alpha)
                return alpha
            alpha *= 0.5  # 缩小步长
        # print("alpha2",alpha)
        return alpha

    def goldstein(self, f, grad, x, d, alpha0=1.0):
        """
        Goldstein 准则

        参数：
        - f：目标函数
        - grad：目标函数的梯度
        - x：当前点
        - d：搜索方向
        - alpha0：初始步长，默认为1.0

        返回：
        - alpha：满足 Goldstein 准则的步长
        """
        alpha = alpha0
        for _ in range(self.max_iters):
            if f(x + alpha * d) <= f(x) + self.c1 * alpha * torch.dot(grad, d) and \
               f(x + alpha * d) >= f(x) + (1 - self.c1) * alpha * torch.dot(grad, d):
                return alpha
            alpha *= 0.5  # 缩小步长
        return alpha

    def wolfe(self, f, grad, x, d, alpha0=1.0):
        """
        Wolfe 准则

        参数：
        - f：目标函数
        - grad：目标函数的梯度
        - x：当前点
        - d：搜索方向
        - alpha0：初始步长，默认为1.0

        返回：
        - alpha：满足 Wolfe 准则的步长
        """
        alpha = alpha0
        for _ in range(self.max_iters):
            if f(x + alpha * d) <= f(x) + self.c1 * alpha * torch.dot(grad, d) and \
               torch.abs(torch.dot(grad, (x + alpha * d - x))) <= -self.c2 * torch.dot(grad, d):
                return alpha
            alpha *= 0.5  # 缩小步长
        return alpha


###1.参数文件
###2.梯度下降算法类
###3.优化函数应该可以是任意可微的

'''
优化参数：学习率（或者选择非精确搜索规则），最大下降步数，最大非精确搜索步数
1.定义优化函数func
2.调用梯度下降decent(func)，梯度下降参数（params）
3.绘制loss下降曲线'''